#include <dialects/Memory/Ops.h>
#include <dialects/PL/PL.h>
#include <dialects/PL/Translation/ToMemory.h>
#include <lcir/Dag2Dag/D2DDriver.h>
#include <lcir/Dag2Dag/Pattern.h>
#include <lcir/Pass/Pass.h>

#include <cassert>

#include "dialects/Builtin/Types.h"
#include "dialects/CFG/TyAndProto.h"
#include "dialects/Func/Func.h"
#include "dialects/PL/Ops.h"
#include "lcir/IR/Argument.h"
#include "lcir/IR/IRBuilder.h"
#include "lcir/IR/IRRewritter.h"
#include "lcir/IR/Lambda.h"
#include "lcir/IR/ValuePrinter.h"
#include "lcir/Utils/Result.h"

namespace lcir::pl {
struct DeclPattern : public OpPattern<DeclVar> {
 public:
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    auto var = expr->getResult();
    auto alloca = b.createAndGetRet<memory::Alloca>({}, var->getType(), var->getInfo());
    alloca->setInfo(var->getInfo());
    for (auto &user : IncRange(var->getUsers())) {
      auto userExpr = user.getExpr();
      IRBuilder b(userExpr);
      if (userExpr->isOp<Assign>() && userExpr->getOperandValue(1) == var) {
        auto store = b.createAndGetRet<memory::Store>({userExpr->getOperandValue(0), alloca});
        continue;
      }
      auto load = b.createAndGetRet<memory::Load>({alloca});
      user.setValue(load);
    }
    b.replaceValue(expr->getResult(), alloca);
    return success();
  }
};
// 处理数组
struct ArrDeclPattern : public OpPattern<DeclArr> {
  Result rewrite(Expression *expr, IRRewritter &b) const override {
    auto var = expr->getResult();
    auto arrDeclOp = expr->getOp<pl::DeclArr>();
    auto size = arrDeclOp->getElementNums();
    auto alloca = b.createAndGetRet<memory::Alloca>({}, var->getType(), var->getInfo(), size);
    for (auto &user : IncRange(var->getUsers())) {
      auto userExpr = user.getExpr();
      IRBuilder b(userExpr);
      auto *load = userExpr->getOp<pl::ArrLoad>();
      if (!load) continue;
      // auto store = b.createAndGetRet<memory::Store>({userExpr->getOperandValue(0), alloca});
      std::vector<Value *> operands{alloca};
      for (size_t i = 0; i < userExpr->getOperands().size(); i++) {
        if (!i) continue;
        operands.push_back(userExpr->getOperandValue(i));
      }
      // std::copy(userExpr->getOperandValues().begin() + 1, userExpr->getOperandValues().end(),
      // std::back_inserter(operands));
      auto ptr = b.createAndGetRet<memory::OffsetPtr>(operands);
      // 对arrayLoad做操作 使用转load，assing转store
      for (auto &user : IncRange(load->getResult()->getUsers())) {
        auto userExpr = user.getExpr();
        if (userExpr->isOp<Assign>() && userExpr->getOperandValue(1) == load->getResult()) {
          auto store = b.createAndGetRet<memory::Store>({userExpr->getOperandValue(0), ptr});
          continue;
        }
        auto load = b.createAndGetRet<memory::Load>({ptr});
        user.setValue(load);
      }
      userExpr->getResult()->rauw(ptr);
    }
    b.replaceValue(expr->getResult(), alloca);
    return success();
  }
};

// 处理数组退化成指针的情况
static void handleArrFuncArgs(Value *arg, IRBuilder &b) {
  auto *argAlloc = b.createAndGetRet<func::ArgAlloca>({}, arg->getType(), arg->getInfo());
  arg->rauw(argAlloc);
  for (auto &user : IncRange(argAlloc->getUsers())) {
    auto userExpr = user.getExpr();
    IRRewritter b(userExpr);
    if (auto *load = userExpr->getOp<pl::ArrLoad>()) {
      // 1.计算偏移量
      std::vector<Value *> operands{argAlloc};
      for (size_t i = 0; i < userExpr->getOperands().size(); i++) {
        if (!i) continue;
        operands.push_back(userExpr->getOperandValue(i));
      }
      auto *ptr = b.createAndGetRet<memory::OffsetPtr>(operands);
      auto argAssign = b.createAndGetRet<func::ArgAssign>({userExpr->getOperandValue(0), argAlloc});
      for (auto &user : IncRange(load->getResult()->getUsers())) {
        auto userExpr = user.getExpr();
        if (userExpr->isOp<Assign>() && userExpr->getOperandValue(1) == load->getResult()) {
          auto store = b.createAndGetRet<memory::Store>({userExpr->getOperandValue(0), ptr});
          continue;
        }
        auto load = b.createAndGetRet<memory::Load>({ptr});
        user.setValue(load);
      }
      userExpr->getResult()->rauw(ptr);
      b.replaceValue(userExpr->getResult(), argAssign);
      continue;
    }
  }
}

// 处理普通情况
static void handleNoramlFuncArgs(Value *arg, IRBuilder &b) {
  auto *argAlloc = b.createAndGetRet<func::ArgAlloca>({}, arg->getType(), arg->getInfo());
  arg->rauw(argAlloc);
  for (auto &user : IncRange(argAlloc->getUsers())) {
    auto userExpr = user.getExpr();
    IRRewritter b(userExpr);
    if (userExpr->isOp<Assign>()) {
      auto argAssign = b.createAndGetRet<func::ArgAssign>({userExpr->getOperandValue(0), argAlloc});

      userExpr->eraseFromLambda();
      b.replaceValue(userExpr->getResult(), argAssign);
      continue;
    }
  }
}

static void handleFuncArgs(Lambda *funcLambda, Lambda *entryLambda) {
  IRBuilder b(*entryLambda->begin());
  for (auto *arg : funcLambda->getArgs()) {
    if(arg->getType()->isa<builtin::Pointer>()) {
      handleArrFuncArgs(arg, b);
    }
    else {
      handleNoramlFuncArgs(arg, b);
    }
  }
}

void ToMemory::run() {
  auto *module = getLambda();

  module->walkOnOp<func::Func>([&](Expression *expr) {
    func::FuncTemp temp(expr);
    auto *entry = temp.getFuncBody()->begin()->getResultProto()->dyn_cast<cfg::BlockProto>();
    assert(entry && "func entry is not a block");
    handleFuncArgs(expr->getOperandValue(0)->dyn_cast<Lambda>(), entry->getLambda());
    return walkSkip();
  });
  D2DDriver driver(module);
  driver.addPatterns<DeclPattern, ArrDeclPattern>();
  driver.run();
  module->walkOnOp<pl::Assign>([&](Expression *expr) {
    std::cerr << expr->print() << "\n";
    expr->eraseFromLambda();
    return walkAdvance();
  });
}
}  // namespace lcir::pl